import datetime
import pandas as pd
import tushare as ts
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from pandas.plotting import register_matplotlib_converters

import QUANTAXIS as QA


register_matplotlib_converters()


def main(codes, start, end):
    # print(codes.head())
    data = QA.QA_fetch_stock_day_adv(codes['code'].to_list(), start, end)
    data = data.to_qfq()
    ind = data.add_func(QA.QA_indicator_beili)

    # data = ind.xs('002044', level='code', axis=0, drop_level=True)
    # figure = plt.figure()
    # axes1 = figure.add_subplot(411)
    # axes2 = figure.add_subplot(412)
    # axes3 = figure.add_subplot(413)
    # axes4 = figure.add_subplot(414)
    # axes1.plot(data.index, data['CLOSE'])
    # axes2.plot(data.index, data[['DIFF', 'DEA']])
    # axes3.bar(data.index, data['MACD'])
    # axes4.plot(data.index, data[['H_BL', 'L_BL']])
    #
    # left, right = axes3.get_xlim()
    # axes3.hlines(y=0, xmin=left, xmax=right, linestyles='dashed')
    # multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)
    # plt.show()

    for _, code in codes.iterrows():
        df = data.data.xs(code['code'], level='code', axis=0, drop_level=True)
        df = QA.QA_indicator_beili(df, valid_bars=5, macd_limit=0.05).fillna(0)
        ind_last_last_bar = df.iloc[-3]
        ind_last_bar = df.iloc[-2]
        ind_current_bar = df.iloc[-1]
        hbl = ind_current_bar.H_BL + ind_last_bar.H_BL + ind_last_last_bar.H_BL
        lbl = ind_current_bar.L_BL + ind_last_bar.L_BL + ind_last_last_bar.L_BL
        if hbl > 0:
            print('code：{}---{}---顶背离---时间 {}---CLOSE {}---LEVEL {}'.format(code['code'], code['name'],
                ind_current_bar.name, ind_current_bar.CLOSE, hbl))
            if hbl == 3:
                plot(df)
        if lbl > 0:
            print('code：{}---{}---底背离---时间 {}---CLOSE {}---LEVEL {}'.format(code['code'], code['name'],
                ind_current_bar.name, ind_current_bar.CLOSE, lbl))
            if lbl == 3:
                plot(df)


def plot(data):
    figure = plt.figure()
    axes1 = figure.add_subplot(411)
    axes2 = figure.add_subplot(412)
    axes3 = figure.add_subplot(413)
    axes4 = figure.add_subplot(414)
    axes1.plot(data.index, data['CLOSE'])
    axes2.plot(data.index, data[['DIFF', 'DEA']])
    axes3.bar(data.index, data['MACD'])
    axes4.plot(data.index, data[['H_BL', 'L_BL']])

    left, right = axes3.get_xlim()
    axes3.hlines(y=0, xmin=left, xmax=right, linestyles='dashed')
    multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)
    plt.show()

if __name__ == '__main__':
    ts_token = '17056d23a59ab71cb979c6a30185e092aba605c4544dac900a3eb7f8'
    ts.set_token(ts_token)
    pro = ts.pro_api()
    data = pro.stock_basic(exchange='', list_status='L', fields='symbol,name,area,industry,list_date')
    data = data[data.symbol.str.startswith('00') | data.symbol.str.startswith('60') | data.symbol.str.startswith('30')]
    # data = data[data.symbol.str.startswith('002044')]
    data = data[data.list_date < '20190101']
    all = data.rename(columns={'symbol': 'code'})
    codes = all[~all.name.str.contains('ST')]
    # codes = codes.iloc[0: 1]
    # codes = pd.read_csv('D:\\PythonPro\\QUANTAXIS\\EXAMPLE\\AI\\data\\hs300.csv', dtype=str)
    # hs300 = ts.get_hs300s()
    # codes = codes['code'].to_list()
    # codes = ['002027']
    # print(codes)
    start = '2020-08-01'
    end = '2021-01-29'
    main(codes, start, end)
